#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <pthread.h>
#include <openssl/md5.h> // 用于计算文件的MD5

#define PORT 8080
#define BUFFER_SIZE 10240
#define MD5_LENGTH 16 // MD5 值的字节长度
#define CHECKSUM_SIZE 4  // 校验和大小，这里使用 4 字节
#define THREAD_COUNT 4

typedef struct {
    int client_sock;
    char filename[256];
    size_t file_size;
    unsigned char md5_expected[MD5_LENGTH]; // 预期的 MD5 值
} ClientData;

// 计算校验和
unsigned int calculate_checksum(const char *data, size_t size) {
    unsigned int checksum = 0;
    for (size_t i = 0; i < size; i++) {
        checksum += (unsigned char)data[i];
    }
    return checksum;
}

void calculate_md5(const char *filename, unsigned char *md5_result) {
    MD5_CTX ctx;
    MD5_Init(&ctx);

    int fd = open(filename, O_RDONLY);
    if (fd < 0) {
        perror("Failed to open file for MD5 calculation");
        return;
    }

    unsigned char buffer[1024];
    ssize_t bytes_read;
    while ((bytes_read = read(fd, buffer, sizeof(buffer))) > 0) {
        MD5_Update(&ctx, buffer, bytes_read);
    }

    close(fd);
    MD5_Final(md5_result, &ctx);
}

void *handle_client(void *arg) {
    ClientData *client_data = (ClientData *)arg;
    char buffer[BUFFER_SIZE];
    size_t total_received = 0;

    // 接收数据并写入文件
    FILE *output_file = fopen(client_data->filename, "wb");
    if (output_file == NULL) {
        perror("File opening failed");
        close(client_data->client_sock);
        free(client_data);
        pthread_exit(NULL);
    }

    char *received_blocks[THREAD_COUNT] = {NULL};
    size_t block_sizes[THREAD_COUNT];
    unsigned int expected_checksums[THREAD_COUNT];

    // 接收每个线程的数据块和校验和
    size_t block_size = client_data->file_size / THREAD_COUNT;
    for (int i = 0; i < THREAD_COUNT; i++) {
        int thread_id;
        unsigned int expected_checksum;

        // 接收线程 ID
        if (recv(client_data->client_sock, &thread_id, sizeof(int), 0) < 0) {
            perror("recv thread ID failed");
            fclose(output_file);
            close(client_data->client_sock);
            free(client_data);
            pthread_exit(NULL);
        }

        // 接收校验和
        if (recv(client_data->client_sock, &expected_checksum, CHECKSUM_SIZE, 0) < 0) {
            perror("recv checksum failed");
            fclose(output_file);
            close(client_data->client_sock);
            free(client_data);
            pthread_exit(NULL);
        }

        expected_checksums[thread_id] = expected_checksum;

        // 接收数据块
        block_sizes[thread_id] = (thread_id == THREAD_COUNT - 1) ? (client_data->file_size - total_received) : block_size;
        received_blocks[thread_id] = (char *)malloc(block_sizes[thread_id]);
        if (received_blocks[thread_id] == NULL) {
            perror("malloc failed");
            fclose(output_file);
            close(client_data->client_sock);
            free(client_data);
            pthread_exit(NULL);
        }

        size_t received_this_block = 0;
        while (received_this_block < block_sizes[thread_id]) {
            ssize_t bytes_received = recv(client_data->client_sock, received_blocks[thread_id] + received_this_block, block_sizes[thread_id] - received_this_block, 0);
            if (bytes_received <= 0) {
                break; // 连接关闭或出错
            }
            received_this_block += bytes_received;
        }

        total_received += block_sizes[thread_id];
    }

    // 按顺序计算校验和并写入文件
    for (int i = 0; i < THREAD_COUNT; i++) {
        unsigned int actual_checksum = calculate_checksum(received_blocks[i], block_sizes[i]);
        if (actual_checksum != expected_checksums[i]) {
            printf("Checksum mismatch for block %d! Expected: %u, Actual: %u\n", i, expected_checksums[i], actual_checksum);
        }

        fwrite(received_blocks[i], 1, block_sizes[i], output_file);
        free(received_blocks[i]);
    }

    // 确保文件数据完全写入磁盘
    fflush(output_file);
    fsync(fileno(output_file));
    fclose(output_file);

    // 计算接收到文件的 MD5
    unsigned char md5_received[MD5_LENGTH];
    calculate_md5(client_data->filename, md5_received);

    // 比较 MD5 值
    if (memcmp(md5_received, client_data->md5_expected, MD5_LENGTH) == 0) {
        printf("File '%s' received successfully with size %zu bytes! MD5 matches.\n", client_data->filename, total_received);
    } else {
        printf("File '%s' received, but MD5 mismatch! Received MD5: ", client_data->filename);
        for (int i = 0; i < MD5_LENGTH; i++) {
            printf("%02x", md5_received[i]);
        }
        printf("\nExpected MD5: ");
        for (int i = 0; i < MD5_LENGTH; i++) {
            printf("%02x", client_data->md5_expected[i]);
        }
        printf("\n");
    }

    // 清理
    close(client_data->client_sock);
    free(client_data);
    pthread_exit(NULL);
}

int main() {
    int sockfd, new_sock;
    struct sockaddr_in server_addr, client_addr;
    socklen_t addr_size;

    // 创建套接字
    sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sockfd < 0) {
        perror("Socket creation failed");
        exit(EXIT_FAILURE);
    }

    // 设置套接字选项以重用地址
    int opt = 1;
    if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        perror("setsockopt failed");
        close(sockfd);
        exit(EXIT_FAILURE);
    }

    // 绑定套接字
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(PORT);
    if (bind(sockfd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        perror("Bind failed");
        close(sockfd);
        exit(EXIT_FAILURE);
    }

    // 监听连接
    if (listen(sockfd, 5) < 0) {
        perror("Listen failed");
        close(sockfd);
        exit(EXIT_FAILURE);
    }

    printf("Waiting for connections...\n");

    while (1) {
        // 接受连接
        ClientData *client_data = malloc(sizeof(ClientData)); // 为每个客户端分配空间
        addr_size = sizeof(client_addr);
        new_sock = accept(sockfd, (struct sockaddr*)&client_addr, &addr_size);
        if (new_sock < 0) {
            perror("Accept failed");
            free(client_data);
            continue; // 继续等待下一个连接
        }

        // 接收文件名长度
        size_t filename_len;
        if (recv(new_sock, &filename_len, sizeof(filename_len), 0) < 0) {
            perror("recv filename length failed");
            close(new_sock);
            free(client_data);
            continue;
        }

        // 接收文件名
        if (recv(new_sock, client_data->filename, filename_len, 0) < 0) {
            perror("recv filename failed");
            close(new_sock);
            free(client_data);
            continue;
        }
        client_data->filename[filename_len] = '\0'; // 确保以 null 结尾

        // 接收文件大小
        if (recv(new_sock, &client_data->file_size, sizeof(client_data->file_size), 0) < 0) {
            perror("recv file size failed");
            close(new_sock);
            free(client_data);
            continue;
        }

        // 接收 MD5 值
        if (recv(new_sock, client_data->md5_expected, MD5_LENGTH, 0) < 0) {
            perror("recv MD5 failed");
            close(new_sock);
            free(client_data);
            continue;
        }

        // 设置客户端 socket
        client_data->client_sock = new_sock;

        // 创建新线程来处理客户端
        pthread_t client_thread;
        if (pthread_create(&client_thread, NULL, handle_client, client_data) != 0) {
            perror("pthread_create failed");
            close(new_sock);
            free(client_data);
            continue;
        }

        // 线程分离，避免内存泄漏
        pthread_detach(client_thread);
    }

    // 清理
    close(sockfd);
    return 0;
}